#include "Kernel.h"
#include "Debugger.h"
#include "Hal/IHardware.h"
#include "Hal/Memory.h"
#include "Utils/Spinlock.h"
#include "Memory/MemoryManager.h"
#include "Memory/PhysicalHeap.h"

CPhysicalHeap sPhysicalHeap;
CPhysicalHeap * PhysicalHeap = &sPhysicalHeap;

extern Address bss;
extern Address end;

Result CPhysicalHeap::Initialize( void )
{
	MEMORY_TRACE( "Initializing...\n" );

	fBaseAddress = 0x200000;
	fOffset = 0;
	fSize = 0x800000;
	fFreePages.fFirst = NULL;
	fFreePages.fCount = 0;
	fUsed = 0;
	fAllocated = 0;
	fFreed = 0;

	for( int i = 0; i < PHYSICAL_HEAP_MAX_ZONE_COUNT; i++ )
	{
		fZones[i] = NULL;
		fZoneSizes[i] = 1 << ( i + 3 );
	}

	MEMORY_TRACE_DETAILED( "Zeroing out heap... %p (%i bytes)\n", fBaseAddress, fSize );
	Memory_Set( (void *)fBaseAddress, 0, fSize );

	return RESULT_OK;
}


void * CPhysicalHeap::Allocate( Size size )
{
	fSpinlock.Lock();
	MEMORY_TRACE( "Allocating, size = %d\n", size );

	void * memory = NULL;

	for( int i = 0; i < PHYSICAL_HEAP_MAX_ZONE_COUNT; i++ )
	{
		if( size <= fZoneSizes[i] )
		{
			memory = _AllocateFromZone( i, size );
			break;
		}
	}

	if( memory == NULL )
		memory = _AllocateChunk( size );

	fAllocated += size;
	fUsed += size;

	MEMORY_TRACE( "Allocated, address = %p\n", memory );
	fSpinlock.Unlock();
	return memory;
}


void * CPhysicalHeap::_AllocateFromZone( int index, Size size )
{
	MEMORY_TRACE_DETAILED( "Allocating from zone, size = %d\n", size );

	if( index < 0 || index >= PHYSICAL_HEAP_MAX_ZONE_COUNT )
		return _AllocateChunk( size );

	PhysicalHeapZone_t * zone = fZones[index];

	// this zone is full, don't use it
	if( zone != NULL && zone->fFreeChunks.fCount == 0 )
	{
		zone->fFlags &= ~PHYSICAL_HEAP_ZONE_IN_USE;
		zone = NULL;
	}

	if( zone == NULL )
	{
		zone = (PhysicalHeapZone_t *)_AllocatePages( MEMORY_PAGE_SIZE );
		zone->fSignature = PHYSICAL_HEAP_ZONE_SIGNATURE;
		zone->fBaseAddress = (Address)zone + sizeof(PhysicalHeapZone_t);
		zone->fChunkSize = fZoneSizes[ index ];
		zone->fChunkCount = ( MEMORY_PAGE_SIZE - sizeof(PhysicalHeapZone_t) ) / zone->fChunkSize;
		zone->fFreeChunks.fFirst = NULL;
		zone->fFreeChunks.fCount = 0;
		zone->fFlags |= PHYSICAL_HEAP_ZONE_IN_USE;
#if MEMORY_USE_WALLS
		zone->fWall = PHYSICAL_HEAP_WALL_SIGNATURE;
#endif

		for( uint32 i = 0; i < zone->fChunkCount; i++ )
		{
			PhysicalHeapNode_t * node = (PhysicalHeapNode_t *)( zone->fBaseAddress + zone->fChunkSize * i );
			node->fNext = zone->fFreeChunks.fFirst;
			zone->fFreeChunks.fFirst = node;
			zone->fFreeChunks.fCount++;
		}

		fZones[index] = zone;
	}

	PhysicalHeapNode_t * node = zone->fFreeChunks.fFirst;
	zone->fFreeChunks.fFirst = node->fNext;
	zone->fFreeChunks.fCount--;

	MEMORY_TRACE_DETAILED( "Zone chunk size = %d, chunks left = %d\n", zone->fChunkSize, zone->fFreeChunks.fCount );
	return (void *)node;
}


void * CPhysicalHeap::_AllocateChunk( Size size )
{
	MEMORY_TRACE_DETAILED( "Allocating raw chunk, size = %d\n", size );

	// make sure we have room for the header as well
	size += sizeof(PhysicalHeapChunk_t);

	PhysicalHeapChunk_t * chunk = (PhysicalHeapChunk_t *)_AllocatePages( size );
	chunk->fSignature = PHYSICAL_HEAP_CHUNK_SIGNATURE;
	chunk->fChunkSize = size;
#if MEMORY_USE_WALLS
	chunk->fWall = PHYSICAL_HEAP_WALL_SIGNATURE;
#endif

	return (void *)( (Address)chunk + sizeof(PhysicalHeapChunk_t) );
}


void * CPhysicalHeap::_AllocatePages( Size size )
{
	MEMORY_TRACE_DETAILED( "Allocating pages, size = %d\n", size );

	// make sure the size is rounded up to the next page
	if( !MEMORY_IS_PAGE_ALIGNED( size ) )
		size = MEMORY_PAGE_ALIGN( size + MEMORY_PAGE_SIZE );

	if( size <= MEMORY_PAGE_SIZE && fFreePages.fCount > 0 )
	{
		// grab the first free page, then remove it from the list
		PhysicalHeapNode_t * node = (PhysicalHeapNode_t *)fFreePages.fFirst;
		fFreePages.fFirst = fFreePages.fFirst->fNext;
		fFreePages.fCount--;

		fUsed += size;
		return node;
	}

	// do we have room in the heap?
	if( fOffset + size > fSize )
		PANIC( "Out of heap space." );

	// grab it of the heap and change the ofset
	void * memory = (void *)( fBaseAddress + fOffset );
	fOffset += size;

	return memory;
}


void CPhysicalHeap::Free( void * memory )
{
	MEMORY_TRACE( "Freeing, address = %08x\n", (Address)memory );

	if( memory == NULL )
		return;

	fSpinlock.Lock();

	uint32 signature = *(uint32 *)MEMORY_PAGE_ALIGN( (Address)memory );

	// is this page a zone or a chunk?
	if( signature == PHYSICAL_HEAP_ZONE_SIGNATURE )
		_FreeToZone( memory );

	else if( signature == PHYSICAL_HEAP_CHUNK_SIGNATURE )
		_FreeChunk( memory );

	else
		PANIC( "Tried to free memory that wasn't a zone or chunk." );

	fSpinlock.Unlock();
}


void CPhysicalHeap::_FreeToZone( void * memory )
{
	MEMORY_TRACE_DETAILED( "Freeing to zone, address = %08x\n", (Address)memory );

	PhysicalHeapZone_t * zone = (PhysicalHeapZone_t *)MEMORY_PAGE_ALIGN( (Address)memory );

	if( zone->fSignature != PHYSICAL_HEAP_ZONE_SIGNATURE )
		PANIC( "Tried to free memory to a zone that wasn't part of a zone." );

#if MEMORY_USE_WALLS
	// make sure the memory wasn't altered
	if( zone->fWall != PHYSICAL_HEAP_WALL_SIGNATURE )
		PANIC( "Zone wall has been altered." );
#endif

	// free the chunk and add it to the free list
	PhysicalHeapNode_t * node = (PhysicalHeapNode_t *)memory;
	node->fNext = zone->fFreeChunks.fFirst;
	zone->fFreeChunks.fFirst = node;
	zone->fFreeChunks.fCount++;

	fUsed  -= zone->fChunkSize;
	fFreed += zone->fChunkSize;

	MEMORY_TRACE_DETAILED( "Zone chunk size = %d, chunks left = %d\n", zone->fChunkSize, zone->fFreeChunks.fCount );

	// if this zone has had all its chunks returned and isn't in use
	if( zone->fFreeChunks.fCount == zone->fChunkCount && !( zone->fFlags & PHYSICAL_HEAP_ZONE_IN_USE ) )
	{
		MEMORY_TRACE( "Freeing zone\n" );
		_FreePage( zone );
	}
}


void CPhysicalHeap::_FreeChunk( void * memory )
{
	MEMORY_TRACE_DETAILED( "Freeing raw chunk, address = %08x\n", (Address)memory );

	PhysicalHeapChunk_t * chunk = (PhysicalHeapChunk_t *)MEMORY_PAGE_ALIGN( (Address)memory );

    ASSERT_SZ( chunk->fSignature == PHYSICAL_HEAP_CHUNK_SIGNATURE, "Tried to free a memory chunk that wasn't a chunk." );

#if MEMORY_USE_WALLS
	ASSERT_SZ( chunk->fWall == PHYSICAL_HEAP_WALL_SIGNATURE, "Chunk wall has been altered." );
#endif

	fUsed  -= chunk->fChunkSize;
	fFreed += chunk->fChunkSize;

	// free all pages that this chunk was using
	uint32 pages = chunk->fChunkSize / MEMORY_PAGE_SIZE;
	for( uint32 i = 0; i < pages; i++ )
		_FreePage( (void *)( (Address)chunk + i * MEMORY_PAGE_SIZE ) );
}


void CPhysicalHeap::_FreePage( void * memory )
{
	MEMORY_TRACE_DETAILED( "Freeing page, address = %08x\n", (Address)memory );

	PhysicalHeapNode_t * pageNode = (PhysicalHeapNode_t *)MEMORY_PAGE_ALIGN( (Address)memory );
	pageNode->fNext = fFreePages.fFirst;
	fFreePages.fFirst = pageNode;
	fFreePages.fCount++;

	fUsed -= MEMORY_PAGE_SIZE;
	fFreed += MEMORY_PAGE_SIZE;
}

